# -*- coding: utf-8 -*-
# @Time    : 2020/12/11 11:56
# @Author  : huangwei
# @File    : video_method.py
# @Software: PyCharm
from multiprocessing import Pool
import os, shutil, ffmpeg, cv2, copy, math
import numpy as np
from psnr import *

# 加入水印的标志，用于判断是否加入了水印
FLAG = np.array(
    [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
     1,
     0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
     1, 0, 1, 0, 1])

BLOCK_SHAPE = np.array((8, 8))  # 进行分块的大小
MARKSTRENGTH = 100  # 值越大 鲁棒性越强，但是对原文件影响越大
GAP = 10  # 视频分段时长


def delete_file( filepath ):
    """用于删除中间产生的临时文件"""
    if os.path.exists(filepath):
        os.remove(filepath)
    else:
        print("%s is not exists." % filepath)


def create_dir( dir_path ):
    """创建文件夹"""
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    else:
        print("%s is exists" % dir_path)


def delete_dir( dir_path ):
    """删除文件夹"""
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
    else:
        print("%s is not exists" % dir_path)


def reformat2mp4( input_filepath, output_filepath ):
    """
    将视频转换为 mp4 格式
    :param input_filepath:
    :param output_filepath:
    :return:
    """
    os.system("ffmpeg -y -i {0} {1}".format(input_filepath, output_filepath))


def get_video_param( video_path ):
    """
    返回视频流中的元素
    :param video_path:
    :return: {duration, width, height, codec_name, pix_fmt, fps, audio}
    """
    probe = ffmpeg.probe(video_path)
    duration = float(probe['format']['duration'])

    video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)

    if video_stream is None:
        raise AssertionError("please check your file, this may not be a video!")

    audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)

    if audio_stream is None:
        is_audio_exists = False
    else:
        is_audio_exists = True

    width = video_stream['width']
    height = video_stream['height']
    codec_name = video_stream['codec_name']
    pix_fmt = video_stream['pix_fmt']
    frame_rate = video_stream['avg_frame_rate']
    fps = round(eval(frame_rate), 2)

    return {"duration": duration, "width": width, "height": height, "codec_name": codec_name, "pix_fmt": pix_fmt,
            "fps": fps, "audio": is_audio_exists}


def insert2yuv( input_filepath, output_filepath, watermark_filepath, img_shape ):
    """
    对 yuv 文件的前 50 帧插入水印
    :param input_filepath:
    :param output_filepath:
    :param watermark_filepath:
    :param img_shape:
    :return:
    """
    INSERT_NUM = 50  # 插入水印的帧数
    fp = open(input_filepath, "rb")  # 读取yuv文件
    input_data = fp.read()

    file_length = len(input_data)  # 整个文件字符长度
    y_length = img_shape[0] * img_shape[1]  # 单帧的 y 向量大小
    yuv_length = int(y_length * 3 / 2)  # 单帧的 yuv 向量的大小
    frames_num = int(file_length / yuv_length)  # 总帧数

    frame = bytearray(yuv_length)  # 用于存储单帧的 yuv 数据

    # 用于存储单帧的 y 向量数据
    # 创建一个 img_shape[0] 行 img_shape[1] 列的二维数组
    data = [[0 for x in range(img_shape[1])] for y in range(img_shape[0])]

    # 用于将 y 向量划分成 8*8 的 block
    strides = 4 * np.array([img_shape[1] * BLOCK_SHAPE[0], BLOCK_SHAPE[1], img_shape[1], 1])
    frame_block_shape = (img_shape[0] // BLOCK_SHAPE[0], img_shape[1] // BLOCK_SHAPE[1], BLOCK_SHAPE[0], BLOCK_SHAPE[1])

    # 初始化 block 的 index
    block_index = [(i, j) for i in range(frame_block_shape[0]) for j in range(frame_block_shape[1])]

    # 生成水印 bit 信息和属性 bit 信息，即将要插入的数据转为二进制
    watermark = cv2.imread(watermark_filepath, cv2.IMREAD_GRAYSCALE)

    watermark_size = watermark.shape[0] * watermark.shape[1]
    insert_size = watermark_size + len(FLAG) + 100

    # 如果水印过大，则将其进行缩放到合适的大小
    if insert_size > len(block_index):
        print("最多可嵌入{}kb信息，水印大小{}kb，因此将水印进行缩放".format(len(block_index) / 1024, watermark_size / 1024))
        max_watermark_size = len(block_index) - len(FLAG) - 100
        scale = (max_watermark_size / watermark_size) ** 0.5
        shape0 = int(watermark.shape[0] * scale)
        shape1 = int(watermark.shape[1] * scale)
        watermark_size = shape0 * shape1

        watermark = cv2.resize(watermark, (shape1, shape0))
        print("新的水印大小为{0}*{1}。".format(watermark.shape[0], watermark.shape[1]))

    watermark_bit = watermark.flatten() > 128

    height = format(watermark.shape[0], "b").zfill(10)  # 将 height, width 转为长度为 10 的二进制数据
    width = format(watermark.shape[1], "b").zfill(10)
    watermark_attr_bit = (height + width) * 5

    # 依次取出每一帧，前 50 帧插入水印
    for i in range(frames_num):
        for j in range(yuv_length):
            frame[j] = input_data[i * yuv_length + j]

        # 判断该帧是否插入数据
        if i < INSERT_NUM:
            # 取出该帧的 y 向量数据
            y_index = 0
            for row in range(img_shape[0]):
                for col in range(img_shape[1]):
                    data[row][col] = frame[y_index]
                    y_index += 1

            frame_y = np.array(data)

            # 将 y 向量划分为 8 * 8 的块
            frame_block = np.lib.stride_tricks.as_strided(frame_y.astype(np.float32), frame_block_shape, strides)
            embed_frame = copy.deepcopy(frame_y)

            print("插入数据帧: %d.--------------------------" % i)

            # 对前 80 个块进行标志位嵌入
            for k in range(len(FLAG)):
                # 对 frame_block[k] 进行离散余弦变换获得DCT系数矩阵
                # 运用余数定理实现水印嵌入
                frame_block_dct = cv2.dct(frame_block[block_index[k]])

                if FLAG[k] == 0:
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 10
                else:
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 30

                frame_block[block_index[k]] = cv2.idct(frame_block_dct)

            # 嵌入水印属性信息 100 bits
            for k in range(len(watermark_attr_bit)):

                frame_block_dct = cv2.dct(frame_block[block_index[k + len(FLAG)]])

                if int(watermark_attr_bit[k]):
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 30
                else:
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 10

                frame_block[block_index[k + len(FLAG)]] = cv2.idct(frame_block_dct)

            # 嵌入水印信息
            for k in range(watermark_size):

                frame_block_dct = cv2.dct(frame_block[block_index[k + len(FLAG) + len(watermark_attr_bit)]])

                if watermark_bit[k] == 0:
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 10
                else:
                    frame_block_dct[-1][-1] = math.floor(frame_block_dct[-1][-1] / MARKSTRENGTH) * MARKSTRENGTH + 30

                frame_block[block_index[k + len(FLAG) + len(watermark_attr_bit)]] = cv2.idct(frame_block_dct)

            # 四维转为二维，还少了整除剩下的部分。
            part_frame = np.concatenate(np.concatenate(frame_block, 1), 1)

            # 将插入水印的部分放回原处补齐整除剩下的部分
            embed_frame[:part_frame.shape[0], :part_frame.shape[1]] = part_frame
            embed_frame = np.clip(embed_frame, a_min=0, a_max=255)

            # 写回 yuv
            y_index = 0
            for row in range(img_shape[0]):
                for col in range(img_shape[1]):
                    frame[y_index] = int(embed_frame[row][col])
                    y_index += 1

        with open(output_filepath, "ab+") as wp:
            wp.write(frame)


def gap_insert_watermark( input_filepath, output_filepath, watermark_filepath, video_param ):
    """
    对每一段视频插入水印
    1. 将视频转换为 yuv 格式
    2. 对yuv前 50 帧插入水印
    3. yuv 转 mp4
    :param watermark_filepath:
    :param video_param:
    :param input_filepath:
    :param output_filepath:
    :return:
    """

    (tmp_path, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    tmp_name, _ = os.path.splitext(tmp_filename)  # 分离文件名和后缀
    yuv_path = "{0}/{1}.yuv".format(tmp_path, tmp_name)
    yuv_embed_path = "{0}/{1}_embed.yuv".format(tmp_path, tmp_name)

    os.system("ffmpeg -y -vcodec {0} -i {1} -pix_fmt {2} {3}".format(video_param['codec_name'], input_filepath,
                                                                     video_param['pix_fmt'], yuv_path))

    delete_file(input_filepath)

    insert2yuv(yuv_path, yuv_embed_path, watermark_filepath, (video_param['height'], video_param['width']))

    delete_file(yuv_path)

    os.system(
        "ffmpeg -y -s:v {0}x{1} -pix_fmt {2} -r {3} -i {4} {5}".format(video_param['width'], video_param['height'],
                                                                       video_param['pix_fmt'], video_param['fps'],
                                                                       yuv_embed_path,
                                                                       output_filepath))

    delete_file(yuv_embed_path)


def split_video( input_filepath, output_filepath, watermark_filepath, video_param ):
    """
    将视频进行分段，每一段开启一个进程
    1. 修改为帧内编码使分段更准确
    :param output_filepath:
    :param watermark_filepath:
    :param video_param:
    :param input_filepath:
    :return:
    """
    # 进行帧内编码
    (tmp_path, _) = os.path.split(input_filepath)  # 分离路径和文件名
    intra_filepath = "{0}/intra.mp4".format(tmp_path)
    os.system("ffmpeg -y -i %s -strict -2 -qscale 0 -intra %s" % (input_filepath, intra_filepath))

    # 分段，每一段开启一个进程
    duration = video_param['duration']
    gap_num = int(duration / GAP)

    p = Pool()
    for i in range(gap_num):
        start_time = i * GAP
        gap_path = "{0}/gap{1}.mp4".format(tmp_path, i)
        gap_insert_path = "{0}/gap{1}_insert.mp4".format(tmp_path, i)
        os.system(
            "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3}".format(start_time, GAP, intra_filepath, gap_path))

        print("---------------")
        process = p.apply_async(gap_insert_watermark,
                                args=(gap_path, gap_insert_path, watermark_filepath, video_param,))

    # 不够 一个 GAP 的部分
    start_time = gap_num * GAP
    gap_path = "{0}/gap{1}.mp4".format(tmp_path, gap_num)
    gap_insert_path = "{0}/gap{1}_insert.mp4".format(tmp_path, gap_num)
    os.system(
        "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3}".format(start_time, duration - start_time, intra_filepath,
                                                                      gap_path))
    process = p.apply_async(gap_insert_watermark,
                            args=(gap_path, gap_insert_path, watermark_filepath, video_param,))

    p.close()
    p.join()

    # 写一个 txt 文件用于连接视频
    list_path = "{0}/list.txt".format(tmp_path)
    fp = open(list_path, "a+")
    for i in range(gap_num + 1):
        fp.write("file gap{0}_insert.mp4\n".format(i))

    fp.close()

    # 视频合并
    os.system("ffmpeg -y -f concat -safe 0 -i {0} -c copy {1}".format(list_path, output_filepath))


def find_flag( input_filepath, img_shape, output_filepath ):
    fp = open(input_filepath, "rb")  # 读取yuv文件
    input_data = fp.read()

    file_length = len(input_data)  # 整个文件字符长度
    y_length = img_shape[0] * img_shape[1]  # 单帧的 y 向量大小
    yuv_length = int(y_length * 3 / 2)  # 单帧的 yuv 向量的大小
    frames_num = int(file_length / yuv_length)  # 总帧数

    frame = bytearray(yuv_length)  # 用于存储单帧的 yuv 数据

    # 用于存储单帧的 y 向量数据
    # 创建一个 img_shape[0] 行 img_shape[1] 列的二维数组
    data = [[0 for x in range(img_shape[1])] for y in range(img_shape[0])]

    # 用于将 y 向量划分成 8*8 的 block
    strides = 4 * np.array([img_shape[1] * BLOCK_SHAPE[0], BLOCK_SHAPE[1], img_shape[1], 1])
    frame_block_shape = (img_shape[0] // BLOCK_SHAPE[0], img_shape[1] // BLOCK_SHAPE[1], BLOCK_SHAPE[0], BLOCK_SHAPE[1])

    # 初始化 block 的 index
    block_index = [(i, j) for i in range(frame_block_shape[0]) for j in range(frame_block_shape[1])]

    # 用来保存psnr值最大的时候的位置
    psnr_max = 0
    max_path = ""

    # 读取每一帧的数据
    for i in range(frames_num):
        for j in range(yuv_length):
            frame[j] = input_data[i * yuv_length + j]

        # 取出该帧的 y 向量数据
        y_index = 0
        for row in range(img_shape[0]):
            for col in range(img_shape[1]):
                data[row][col] = frame[y_index]
                y_index += 1

        frame_y = np.array(data)

        # 将 y 向量划分为 8 * 8 的块
        frame_block = np.lib.stride_tricks.as_strided(frame_y.astype(np.float32), frame_block_shape, strides)

        # 提取出前80 个 block 中插入的数据
        temp_list = []
        for k in range(80):
            frame_block_dct = cv2.dct(frame_block[block_index[k]])
            temp_list.append(frame_block_dct[-1][-1] % MARKSTRENGTH)
            frame_block[block_index[k]] = cv2.idct(frame_block_dct)

        list_bit = np.array(temp_list) > 15
        print(list_bit)

        # 将80 转成 16 去除误差
        temp_bit = []
        for k in range(16):
            temp = int(list_bit[k]) + int(list_bit[k + 16]) + int(list_bit[k + 16 * 2]) + int(
                list_bit[k + 16 * 3]) + int(list_bit[k + 16 * 4])
            temp_bit.append(temp)

        flag_bit = np.array(temp_bit) > 2

        count_start = 0
        for k in range(16):
            if flag_bit[k] == FLAG[k]:
                count_start += 1

        if count_start > 13:
            # 则说明该帧存在水印，提取出来
            # 先提取 81 到 180 个 block 的水印属性信息
            temp_list = []
            for k in range(80, 180):
                frame_block_dct = cv2.dct(frame_block[block_index[k]])
                temp_list.append(frame_block_dct[-1][-1] % MARKSTRENGTH)
                frame_block[block_index[k]] = cv2.idct(frame_block_dct)

            list_bit = np.array(temp_list) > 15

            height_arr = []
            width_arr = []
            for k in range(5):
                h_tmp = list_bit[0 + k * 20] * 512 + list_bit[1 + k * 20] * 256 + list_bit[2 + k * 20] * 128 + list_bit[
                    3 + k * 20] * 64 + list_bit[4 + k * 20] * 32 + list_bit[5 + k * 20] * 16 + list_bit[
                            6 + k * 20] * 8 + \
                        list_bit[7 + k * 20] * 4 + list_bit[8 + k * 20] * 2 + list_bit[9 + k * 20]
                height_arr.append(h_tmp)

                w_tmp = list_bit[10 + k * 20] * 512 + list_bit[11 + k * 20] * 256 + list_bit[12 + k * 20] * 128 + \
                        list_bit[
                            13 + k * 20] * 64 + list_bit[14 + k * 20] * 32 + list_bit[15 + k * 20] * 16 + list_bit[
                            16 + k * 20] * 8 + list_bit[17 + k * 20] * 4 + list_bit[18 + k * 20] * 2 + list_bit[
                            19 + k * 20]
                width_arr.append(w_tmp)

            height = np.argmax(np.bincount(height_arr))
            width = np.argmax(np.bincount(width_arr))
            print("height", height)
            print("width:", width)

            if width < 250 and height < 100:
                # 提取水印信息
                watermark_size = width * height

                temp_list = []
                for k in range(180, 180 + watermark_size):
                    frame_block_dct = cv2.dct(frame_block[block_index[k]])
                    temp_list.append(frame_block_dct[-1][-1] % MARKSTRENGTH)
                    frame_block[block_index[k]] = cv2.idct(frame_block_dct)

                list_bit = np.array(temp_list) > 15

                watermark_data = [[0 for x in range(width)] for y in range(height)]

                index = 0
                for row in range(height):
                    for col in range(width):
                        watermark_data[row][col] = int(list_bit[index])
                        index += 1

                watermark = 255 * np.array(watermark_data)

                psnr = water_psnr(watermark)
                # 如果psnr > 15 直接保存返回
                # 否则找出最大的 psnr 时保存返回
                if psnr > 15:
                    print("视频中存在水印，将水印提取保存在 %s" % output_filepath)
                    cv2.imwrite(output_filepath, watermark)
                    return True
                elif psnr > psnr_max:
                    psnr_max = psnr
                    cv2.imwrite(output_filepath, watermark)  # 覆盖之前的

    if os.path.exists(output_filepath):
        print("视频中存在水印，将水印提取保存在 %s" % output_filepath)
        return True

    print("视频中不存在水印！")
    return False


def compress_video( input_filepath, compress_output_filepath ):
    """
    对视频进行压缩，将视频转换成 800*400 或 400*800 分辨率的视频
    :param input_filepath:
    :param compress_output_filepath:
    :return:
    """
    video_param = get_video_param(input_filepath)
    width = video_param['width']
    height = video_param['height']

    # 需要是偶数
    if width >= height:
        new_width = 800
        new_height = 400
    else:
        new_height = 800
        new_width = 400

    os.system("ffmpeg -y -i {0} -s {1}x{2} {3}".format(input_filepath, new_width, new_height, compress_output_filepath))
